/*
 * Copyright (C) 2001-2011 Jacek Sieka, arnetheduck on gmail point com
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

#ifndef DCPLUSPLUS_DCPP_MERKLE_TREE_H
#define DCPLUSPLUS_DCPP_MERKLE_TREE_H

#include "debug.h"
#include "typedefs.h"
#include "TigerHash.h"
#include "Encoder.h"
#include "HashValue.h"

const int64_t MIN_BLOCK_SIZE = 65536;

/**
 * A class that represents a Merkle Tree hash. Storing
 * only the leaves of the tree, it is rather memory efficient,
 * but can still take a significant amount of memory during / after
 * hash generation.
 * The root hash produced can be used like any
 * other hash to verify the integrity of a whole file, while
 * the leaves provide checking of smaller parts of the file.
 */
template < class Hasher, size_t baseBlockSize = 1024 >
class MerkleTree
{
	public:
		static const size_t BITS = Hasher::BITS;
		static const size_t BYTES = Hasher::BYTES;
		static const size_t BASE_BLOCK_SIZE = baseBlockSize;
		
		typedef HashValue<Hasher> MerkleValue;
		typedef vector<MerkleValue> MerkleList;
		typedef typename MerkleList::const_iterator MerkleIter;
		
		MerkleTree() : fileSize(0), blockSize(baseBlockSize) { }
		MerkleTree(int64_t aBlockSize) : fileSize(0), blockSize(aBlockSize) { }
		
		/**
		 * Loads a set of leaf hashes, calculating the root
		 * @param data Pointer to (aFileSize + aBlockSize - 1) / aBlockSize) hash values,
		 *             stored consecutively left to right
		 */
		MerkleTree(int64_t aFileSize, int64_t aBlockSize, uint8_t* aData) :
			fileSize(aFileSize), blockSize(aBlockSize)
		{
			const size_t n = calcBlocks(aFileSize, aBlockSize);
			leaves.reserve(n); //[+]PPA
			for (size_t i = 0; i < n; i++)
				leaves.push_back(MerkleValue(aData + i * Hasher::BYTES));
			calcRoot();
		}
		
		/** Initialise a single root tree */
		MerkleTree(int64_t aFileSize, int64_t aBlockSize, const MerkleValue& aRoot) : root(aRoot), fileSize(aFileSize), blockSize(aBlockSize)
		{
			leaves.push_back(root);
		}
		
		~MerkleTree()
		{
		}
		
		static int64_t calcBlockSize(int64_t aFileSize, int maxLevels)
		{
			int64_t tmp = baseBlockSize;
			int64_t maxHashes = ((int64_t)1) << (maxLevels - 1);
			while ((maxHashes * tmp) < aFileSize)
				tmp *= 2;
			return tmp;
		}
		
		static size_t calcBlocks(int64_t aFileSize, int64_t aBlockSize)
		{
			return max((size_t)((aFileSize + aBlockSize - 1) / aBlockSize), (size_t)1);
		}
		static uint16_t calcBlocks(int64_t aFileSize)
		{
			return (uint16_t)calcBlocks(aFileSize, calcBlockSize(aFileSize, 10));
		}
		
		/**
		 * Update the merkle tree.
		 * @param len Length of data, must be a multiple of BASE_BLOCK_SIZE, unless it's
		 *            the last block.
		 */
		void update(const void* data, size_t len)
		{
			uint8_t* buf = (uint8_t*)data;
			uint8_t zero = 0;
			size_t i = 0;
			
			// Skip empty data sets if we already added at least one of them...
			if (len == 0 && !(leaves.empty() && blocks.empty()))
				return;
				
			do
			{
				size_t n = min(size_t(BASE_BLOCK_SIZE), len - i);
				Hasher h;
				h.update(&zero, 1);
				h.update(buf + i, n);
				if ((int64_t)baseBlockSize < blockSize)
				{
					blocks.push_back(make_pair(MerkleValue(h.finalize()), baseBlockSize));
					reduceBlocks();
				}
				else
				{
					leaves.push_back(MerkleValue(h.finalize()));
				}
				i += n;
			}
			while (i < len);
			fileSize += len;
		}
		
		uint8_t* finalize()
		{
			// No updates yet, make sure we have at least one leaf for 0-length files...
			if (leaves.empty() && blocks.empty())
			{
				update(0, 0);
			}
			while (blocks.size() > 1)
			{
				MerkleBlock& a = blocks[blocks.size() - 2];
				MerkleBlock& b = blocks[blocks.size() - 1];
				a.first = combine(a.first, b.first);
				blocks.pop_back();
			}
			dcassert(blocks.size() == 0 || blocks.size() == 1);
			if (!blocks.empty())
			{
				leaves.push_back(blocks[0].first);
			}
			calcRoot();
			return root.data;
		}
		
		MerkleValue& getRoot()
		{
			return root;
		}
		const MerkleValue& getRoot() const
		{
			return root;
		}
		MerkleList& getLeaves()
		{
			return leaves;
		}
		const MerkleList& getLeaves() const
		{
			return leaves;
		}
		
		int64_t getBlockSize() const
		{
			return blockSize;
		}
		void setBlockSize(int64_t aSize)
		{
			blockSize = aSize;
		}
		
		int64_t getFileSize() const
		{
			return fileSize;
		}
		void setFileSize(int64_t aSize)
		{
			fileSize = aSize;
		}
		
		bool verifyRoot(const uint8_t* aRoot)
		{
			return memcmp(aRoot, getRoot().data(), BYTES) == 0;
		}
		
		void calcRoot()
		{
			root = getHash(0, fileSize);
		}
		
		ByteVector getLeafData()
		{
			ByteVector buf(getLeaves().size() * BYTES);
			const size_t l_size = getLeaves().size(); //[+]PPA
			uint8_t* p = &buf[0];
			for (size_t i = 0; i < l_size; ++i)
			{
				memcpy(p + i * BYTES, &getLeaves()[i], BYTES);
			}
			return buf;
		}
		
	private:
		typedef pair<MerkleValue, int64_t> MerkleBlock;
		typedef vector<MerkleBlock> MBList;
		
		MBList blocks;
		
		MerkleList leaves;
		
		MerkleValue root;
		/** Total size of hashed data */
		int64_t fileSize;
		/** Final block size */
		int64_t blockSize;
		
		MerkleValue getHash(int64_t start, int64_t length)
		{
			dcassert((start % blockSize) == 0);
			if (length <= blockSize)
			{
				start /= blockSize;
				dcassert(start < (int64_t)leaves.size());
				if (start < static_cast<int64_t>(leaves.size())) //[+]PPA
					return leaves[static_cast<size_t>(start)];
				else
					return MerkleValue(); // TODO - raise exception
					
					
			}
			else
			{
				int64_t l = blockSize;
				while (l * 2 < length)
					l *= 2;
				return combine(getHash(start, l), getHash(start + l, length - l));
			}
		}
		
		MerkleValue combine(const MerkleValue& a, const MerkleValue& b)
		{
			uint8_t one = 1;
			Hasher h;
			h.update(&one, 1);
			h.update(a.data, MerkleValue::BYTES);
			h.update(b.data, MerkleValue::BYTES);
			return MerkleValue(h.finalize());
		}
		
		void reduceBlocks()
		{
			while (blocks.size() > 1)
			{
				MerkleBlock& a = blocks[blocks.size() - 2];
				MerkleBlock& b = blocks[blocks.size() - 1];
				if (a.second == b.second)
				{
					if (a.second * 2 == blockSize)
					{
						leaves.push_back(combine(a.first, b.first));
						blocks.pop_back();
						blocks.pop_back();
					}
					else
					{
						a.second *= 2;
						a.first = combine(a.first, b.first);
						blocks.pop_back();
					}
				}
				else
				{
					break;
				}
			}
		}
};

typedef MerkleTree<TigerHash> TigerTree;
typedef TigerTree::MerkleValue TTHValue;

struct TTFilter
{
		TTFilter() : tt(1024 * 1024 * 1024) { }
		void operator()(const void* data, size_t len)
		{
			tt.update(data, len);
		}
		inline TigerTree& getTree()
		{
			return tt;
		}
	private:
		TigerTree tt;
};

#endif // DCPLUSPLUS_DCPP_MERKLE_TREE_H

/**
 * @file
 * $Id: MerkleTree.h 568 2011-07-24 18:28:43Z bigmuscle $
 */
